from pathlib import Path
from pydantic import BaseModel, Field
import os


def _default_embedding_model() -> str:
    # 如果用户没有设置环境变量，使用一个体积较小的开源模型
    env_val = os.environ.get("RAG_EMBEDDING_MODEL")
    if env_val:
        return env_val
    return "C:\\Users\\LenovoTest\\.cache\\modelscope\\hub\\models\\Qwen\\Qwen3-Embedding-4B"


class RAGConfig(BaseModel):
    persist_dir: Path = Field(default=Path(
        os.environ.get("RAG_PERSIST_DIR", "rag_store")))
    embedding_model: str = Field(default_factory=_default_embedding_model)
    collection_name: str = Field(
        default=os.environ.get("RAG_COLLECTION", "documents"))
    chunk_size: int = 2048
    chunk_overlap: int = 100
    top_k: int = 4
    # 向量化批大小，可通过环境变量 RAG_BATCH_SIZE 覆盖
    batch_size: int = Field(default=int(os.environ.get("RAG_BATCH_SIZE", 8)))


config = RAGConfig()
config.persist_dir.mkdir(parents=True, exist_ok=True)
